/*
 * HSM Proxy Project.
 * Copyright (C) 2013 FedICT.
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License version
 * 3.0 as published by the Free Software Foundation.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, see 
 * http://www.gnu.org/licenses/.
 */

package test.integ.be.fedict.hsm.ws;

import java.security.InvalidAlgorithmParameterException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.Collections;
import java.util.Date;
import java.util.GregorianCalendar;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.TimeZone;

import javax.xml.crypto.MarshalException;
import javax.xml.crypto.dom.DOMStructure;
import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureException;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.ExcC14NParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;
import javax.xml.datatype.DatatypeConfigurationException;
import javax.xml.datatype.DatatypeFactory;
import javax.xml.datatype.XMLGregorianCalendar;
import javax.xml.namespace.QName;
import javax.xml.soap.SOAPException;
import javax.xml.soap.SOAPMessage;
import javax.xml.soap.SOAPPart;
import javax.xml.ws.ProtocolException;
import javax.xml.ws.handler.MessageContext;
import javax.xml.ws.handler.soap.SOAPHandler;
import javax.xml.ws.handler.soap.SOAPMessageContext;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.w3c.dom.Attr;
import org.w3c.dom.DOMException;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

public class WSSecurityTestSOAPHandler implements
		SOAPHandler<SOAPMessageContext> {

	private static final Log LOG = LogFactory
			.getLog(WSSecurityTestSOAPHandler.class);

	private static final String WSSE_NAMESPACE = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-secext-1.0.xsd";

	private static final String WSU_NAMESPACE = "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-utility-1.0.xsd";

	private static final String SOAP_NAMESPACE = "http://www.w3.org/2003/05/soap-envelope";

	private static final String XMLNS_NS = "http://www.w3.org/2000/xmlns/";

	private boolean addTimestamp;

	private X509Certificate certificate;

	private PrivateKey privateKey;

	private String digestAlgorithm;

	private String signatureAlgorithm;

	private boolean signBody;

	private boolean signBinarySecurityToken;

	public WSSecurityTestSOAPHandler() {
		super();
		this.digestAlgorithm = DigestMethod.SHA256;
		this.signatureAlgorithm = "http://www.w3.org/2001/04/xmldsig-more#rsa-sha256";
		this.signBody = true;
	}

	public void addTimestamp(boolean addTimestamp) {
		this.addTimestamp = addTimestamp;
	}

	public void setDigestAlgorithm(String digestAlgorithm) {
		this.digestAlgorithm = digestAlgorithm;
	}

	public void setSignatureAlgorithm(String signatureAlgorithm) {
		this.signatureAlgorithm = signatureAlgorithm;
	}

	public void setSignBinarySecurityToken(boolean signBinarySecurityToken) {
		this.signBinarySecurityToken = signBinarySecurityToken;
	}

	@Override
	public boolean handleMessage(SOAPMessageContext context) {
		Boolean outboundProperty = (Boolean) context
				.get(MessageContext.MESSAGE_OUTBOUND_PROPERTY);
		if (true == outboundProperty.booleanValue()) {
			try {
				handleOutboundMessage(context);
			} catch (Exception e) {
				LOG.debug("error adding WS-Security header: " + e.getMessage(),
						e);
				throw new ProtocolException(e);
			}
		}
		return true;
	}

	private void handleOutboundMessage(SOAPMessageContext context)
			throws SOAPException, DatatypeConfigurationException,
			CertificateEncodingException, DOMException,
			NoSuchAlgorithmException, InvalidAlgorithmParameterException,
			MarshalException, XMLSignatureException, NoSuchProviderException {
		SOAPMessage soapMessage = context.getMessage();
		SOAPPart soapPart = soapMessage.getSOAPPart();

		Element soapEnvelopeElement = soapPart.getDocumentElement();
		String soapPrefix = soapEnvelopeElement.getPrefix();
		LOG.debug("SOAP prefix: " + soapPrefix);
		Element soapHeaderElement = soapPart.createElementNS(SOAP_NAMESPACE,
				soapPrefix + ":Header");
		Element soapBodyElement = (Element) soapEnvelopeElement.getFirstChild();
		soapBodyElement.setAttributeNS(XMLNS_NS, "xmlns:wsu", WSU_NAMESPACE);
		soapBodyElement.setAttributeNS(WSU_NAMESPACE, "wsu:Id", "Body");
		soapEnvelopeElement.insertBefore(soapHeaderElement, soapBodyElement);

		LOG.debug("adding WS-Security SOAP header");
		Element wsSecurityHeaderElement = soapPart.createElementNS(
				WSSE_NAMESPACE, "wsse:Security");
		soapHeaderElement.appendChild(wsSecurityHeaderElement);
		wsSecurityHeaderElement.setAttributeNS(XMLNS_NS, "xmlns:wsse",
				WSSE_NAMESPACE);
		wsSecurityHeaderElement.setAttributeNS(XMLNS_NS, "xmlns:wsu",
				WSU_NAMESPACE);
		wsSecurityHeaderElement.setAttributeNS(SOAP_NAMESPACE, soapPrefix
				+ ":mustUnderstand", "true");

		Element tsElement = addTimestamp(wsSecurityHeaderElement);
		addBinarySecurityToken(wsSecurityHeaderElement);
		addSignature(wsSecurityHeaderElement, tsElement, soapBodyElement);
	}

	private void addSignature(Element wsSecurityHeaderElement,
			Element tsElement, Element bodyElement)
			throws NoSuchAlgorithmException,
			InvalidAlgorithmParameterException, MarshalException,
			XMLSignatureException, NoSuchProviderException, SOAPException {
		if (null == this.privateKey) {
			return;
		}
		DOMSignContext domSignContext = new DOMSignContext(this.privateKey,
				wsSecurityHeaderElement);
		domSignContext.setDefaultNamespacePrefix("ds");
		domSignContext.setIdAttributeNS(tsElement, WSU_NAMESPACE, "Id");
		domSignContext.setIdAttributeNS(bodyElement, WSU_NAMESPACE, "Id");
		LOG.debug("Timestamp element found: "
				+ (null != domSignContext.getElementById("TS")));
		XMLSignatureFactory xmlSignatureFactory = XMLSignatureFactory
				.getInstance("DOM");

		List<Reference> references = new LinkedList<Reference>();

		List<String> tsPrefixes = new LinkedList<String>();
		tsPrefixes.add("wsse");
		tsPrefixes.add("S");
		ExcC14NParameterSpec tsTransformSpec = new ExcC14NParameterSpec(
				tsPrefixes);
		Reference tsReference = xmlSignatureFactory
				.newReference("#TS", xmlSignatureFactory.newDigestMethod(
						this.digestAlgorithm, null), Collections
						.singletonList(xmlSignatureFactory.newTransform(
								CanonicalizationMethod.EXCLUSIVE,
								tsTransformSpec)), null, null);
		references.add(tsReference);

		if (this.signBody) {
			List<String> bodyPrefixes = new LinkedList<String>();
			ExcC14NParameterSpec bodyTransformSpec = new ExcC14NParameterSpec(
					bodyPrefixes);
			Reference bodyReference = xmlSignatureFactory.newReference("#Body",
					xmlSignatureFactory.newDigestMethod(this.digestAlgorithm,
							null), Collections
							.singletonList(xmlSignatureFactory.newTransform(
									CanonicalizationMethod.EXCLUSIVE,
									bodyTransformSpec)), null, null);
			references.add(bodyReference);
		}

		if (this.signBinarySecurityToken) {
			Reference bstReference = xmlSignatureFactory
					.newReference("#X509", xmlSignatureFactory.newDigestMethod(
							this.digestAlgorithm, null), Collections
							.singletonList(xmlSignatureFactory.newTransform(
									CanonicalizationMethod.EXCLUSIVE,
									(TransformParameterSpec) null)), null, null);
			references.add(bstReference);
		}

		SignedInfo signedInfo = xmlSignatureFactory.newSignedInfo(
				xmlSignatureFactory.newCanonicalizationMethod(
						CanonicalizationMethod.EXCLUSIVE,
						(C14NMethodParameterSpec) null), xmlSignatureFactory
						.newSignatureMethod(this.signatureAlgorithm, null),
				references);

		KeyInfoFactory keyInfoFactory = xmlSignatureFactory.getKeyInfoFactory();
		Document document = wsSecurityHeaderElement.getOwnerDocument();
		Element securityTokenReferenceElement = document.createElementNS(
				WSSE_NAMESPACE, "wsse:SecurityTokenReference");
		Element referenceElement = document.createElementNS(WSSE_NAMESPACE,
				"wsse:Reference");
		referenceElement
				.setAttribute(
						"ValueType",
						"http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-x509-token-profile-1.0#X509v3");
		referenceElement.setAttribute("URI", "#X509");
		securityTokenReferenceElement.appendChild(referenceElement);
		KeyInfo keyInfo = keyInfoFactory
				.newKeyInfo(Collections.singletonList(new DOMStructure(
						securityTokenReferenceElement)));

		XMLSignature xmlSignature = xmlSignatureFactory.newXMLSignature(
				signedInfo, keyInfo, null, "SIG", null);
		xmlSignature.sign(domSignContext);
	}

	private void addBinarySecurityToken(Element wsSecurityHeaderElement)
			throws SOAPException, CertificateEncodingException, DOMException {
		if (null == this.certificate) {
			return;
		}
		Document document = wsSecurityHeaderElement.getOwnerDocument();

		Element binarySecurityTokenElement = document.createElementNS(
				WSSE_NAMESPACE, "wsse:BinarySecurityToken");
		binarySecurityTokenElement
				.setAttribute(
						"EncodingType",
						"http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-soap-message-security-1.0#Base64Binary");
		binarySecurityTokenElement
				.setAttribute(
						"ValueType",
						"http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-x509-token-profile-1.0#X509v3");
		binarySecurityTokenElement.setAttributeNS(WSU_NAMESPACE, "wsu:Id",
				"X509");
		binarySecurityTokenElement.setTextContent(Base64
				.encodeBase64String(this.certificate.getEncoded()));
		wsSecurityHeaderElement.appendChild(binarySecurityTokenElement);
	}

	private Element addTimestamp(Element wsSecurityHeaderElement)
			throws SOAPException, DatatypeConfigurationException {
		if (false == this.addTimestamp) {
			return null;
		}
		Document document = wsSecurityHeaderElement.getOwnerDocument();
		Element timestampElement = document.createElementNS(WSU_NAMESPACE,
				"wsu:Timestamp");
		timestampElement.setAttributeNS(WSU_NAMESPACE, "wsu:Id", "TS");
		Attr idAttr = timestampElement.getAttributeNodeNS(WSU_NAMESPACE, "Id");
		timestampElement.setIdAttributeNode(idAttr, true);

		Element createdElement = document.createElementNS(WSU_NAMESPACE,
				"wsu:Created");
		DatatypeFactory datatypeFactory = DatatypeFactory.newInstance();
		GregorianCalendar gregorianCalendar = new GregorianCalendar();
		Date now = new Date();
		gregorianCalendar.setTime(now);
		gregorianCalendar.setTimeZone(TimeZone.getTimeZone("UTC"));
		XMLGregorianCalendar xmlGregorianCalendar = datatypeFactory
				.newXMLGregorianCalendar(gregorianCalendar);
		createdElement.setTextContent(xmlGregorianCalendar.toXMLFormat());
		timestampElement.appendChild(createdElement);

		Element expiresElement = document.createElementNS(WSU_NAMESPACE,
				"wsu:Expires");
		Date expiresDate = new Date(now.getTime() + 1000 * 60 * 5);
		gregorianCalendar.setTime(expiresDate);
		xmlGregorianCalendar = datatypeFactory
				.newXMLGregorianCalendar(gregorianCalendar);
		expiresElement.setTextContent(xmlGregorianCalendar.toXMLFormat());
		timestampElement.appendChild(expiresElement);
		wsSecurityHeaderElement.appendChild(timestampElement);
		return timestampElement;
	}

	@Override
	public boolean handleFault(SOAPMessageContext context) {
		return true;
	}

	@Override
	public void close(MessageContext context) {
	}

	@Override
	public Set<QName> getHeaders() {
		return null;
	}

	public void addBinarySecurityToken(X509Certificate certificate) {
		this.certificate = certificate;
	}

	public void addSignature(PrivateKey privateKey) {
		this.privateKey = privateKey;
	}

	public void setSignBody(boolean signBody) {
		this.signBody = signBody;
	}
}
